"""
    Utilities library for evaluation.

Code borrowed from tensorlow-gan library.
We do not claim any ownership on this code and you should refer to the LICENCE
of the tensorflow-gan library.

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import numpy as np
import PIL

import tensorflow as tf
import data_provider
import tensorflow_gan as tfgan  # tf


def get_activations(get_images_fn, num_batches, get_logits=False):
    """Get Inception activations.

    Use TF-GAN utility to avoid holding images or Inception activations in
    memory all at once.

    Args:
      get_images_fn: A function that takes no arguments and returns images.
      num_batches: The number of batches to fetch at a time.
      get_logits: If `True`, return (logits, pools). Otherwise just return pools.

    Returns:
      1 or 2 Tensors of Inception activations.
    """
    inception_img_sz = tfgan.eval.INCEPTION_DEFAULT_IMAGE_SIZE
    outputs = tfgan.eval.sample_and_run_inception(sample_fn=lambda _: tf.compat.v1.image.resize(get_images_fn(),
                                                                                                [inception_img_sz, inception_img_sz],
                                                                                                method=tf.image.ResizeMethod.BILINEAR),
                                                  sample_inputs=[1.0] * num_batches)
    if get_logits:
        return outputs['logits'], outputs['pool_3']
    else:
        return outputs['pool_3']


def get_activations_from_dataset(image_ds, num_batches, get_logits=False):
    """Get Inception activations.

    Args:
      image_ds: tf.Dataset for images.
      num_batches: The number of batches to fetch at a time.
      get_logits: If `True`, return (logits, pools). Otherwise just return pools.

    Returns:
      1 or 2 Tensors of Inception activations.
    """
    iterator = tf.compat.v1.data.make_one_shot_iterator(image_ds)

    get_images_fn = iterator.get_next
    return get_activations(get_images_fn, num_batches, get_logits)


def get_real_activations(batch_size,
                         num_batches,
                         shuffle_buffer_size=100000,
                         split='validation',
                         get_logits=False):
    """Fetches batches inception pools and images.

    NOTE: This function runs inference on an Inception network, so it would be
    more efficient to run this on GPU or TPU than on CPU.

    Args:
      batch_size: The number of elements in a single minibatch.
      num_batches: The number of batches to fetch at a time.
      shuffle_buffer_size: The number of records to load before shuffling. Larger
          means more likely randomization.
      split: Shuffle if 'train', else deterministic.
      get_logits: If `True`, return (logits, pools). Otherwise just return pools.

    Returns:
      A Tensor of `real_pools` or (`real_logits`, `real_pools`) with batch
      dimension (batch_size * num_batches).
    """
    ds = data_provider.provide_dataset(batch_size, shuffle_buffer_size, split)
    ds = ds.map(lambda img, lbl: img)  # Remove labels.
    return get_activations_from_dataset(ds, num_batches, get_logits)


def print_debug_statistics(image, labels, dbg_messge_prefix, on_tpu):
    """Adds a Print directive to an image tensor which prints debug statistics."""
    if on_tpu:
        # Print operations are not supported on TPUs.
        return image, labels
    image_means = tf.reduce_mean(input_tensor=image, axis=0, keepdims=True)
    image_vars = tf.reduce_mean(
        input_tensor=tf.math.squared_difference(image, image_means),
        axis=0,
        keepdims=True)
    image = tf.compat.v1.Print(
        image, [
            tf.reduce_mean(input_tensor=image_means),
            tf.reduce_mean(input_tensor=image_vars)
        ],
        dbg_messge_prefix + ' mean and average var',
        first_n=1)
    labels = tf.compat.v1.Print(
        labels, [labels, labels.shape],
        dbg_messge_prefix + ' sparse labels',
        first_n=2)
    return image, labels


def log_and_summarize_variables(var_list, dbg_messge, on_tpu):
    """Logs given variables, summarizes sigma_ratio_vars."""
    tf.compat.v1.logging.info(dbg_messge + str(var_list))
    sigma_ratio_vars = [var for var in var_list if 'sigma_ratio' in var.name]
    tf.compat.v1.logging.info('sigma_ratio_vars %s %s', dbg_messge,
                              sigma_ratio_vars)
    # Reset the name scope so the summary names are displayed as passed to the
    # summary function.
    if not on_tpu:
        # The TPU estimator doesn't support summaries.
        with tf.compat.v1.name_scope(name=None):
            for var in sigma_ratio_vars:
                tf.compat.v1.summary.scalar(
                    'sigma_ratio_vars/' + var.name, var)


def predict_and_write_images(estimator, input_fn, model_dir, filename_suffix):
    """Generates images and write them to the model dir.

    Args:
      estimator: An object of type tfgan.estimator.GANEstimator or
        tfgan.estimator.TPUGANEstimator for performing the predictions.
      input_fn: An input_fn function to be used by `estimator.predict`.
      model_dir: The model directory (the images will be saved inside an 'images'
        subdirectory).
      filename_suffix: A suffix to append to the image file names.
    """
    # Generate images.
    image_iterator = estimator.predict(input_fn)
    if isinstance(estimator, tfgan.estimator.TPUGANEstimator):
        predictions = np.array(
            [next(image_iterator)['generated_data'] for _ in range(16)])
    else:
        predictions = np.array([next(image_iterator) for _ in range(16)])
    # Write images to disk.
    output_dir = os.path.join(model_dir, 'images')
    if not tf.io.gfile.exists(output_dir):
        tf.io.gfile.makedirs(output_dir)
    # Generate a grid of images and write it to disk.
    image_grid = tfgan.eval.python_image_grid(predictions, grid_shape=(4, 4))
    grid_fname = os.path.join(output_dir, 'grid_%s.png' % filename_suffix)
    _write_image_to_disk(image_grid, grid_fname)


def _write_image_to_disk(image, filename):
    with tf.io.gfile.GFile(filename, 'w') as f:
        # Convert tiled_image from float32 in [-1, 1] to unit8 [0, 255].
        img_np = (255 / 2.0) * (image + 1.0)
        pil_image = PIL.Image.fromarray(img_np.astype(np.uint8))
        pil_image.convert('RGB').save(f, 'PNG')
    tf.compat.v1.logging.info('Wrote output to: %s', filename)
